import chromadb
from chromadb.config import Settings
from openai import OpenAI


class FinancialSituationMemory:
    def __init__(self, name, config):
        if config["backend_url"] == "http://localhost:11434/v1":
            self.embedding = "nomic-embed-text"
        elif "openrouter.ai" in config["backend_url"]:
            # OpenRouter支持的嵌入模型 - 尝试多个选项
            self.embedding = "openai/text-embedding-ada-002"
        else:
            self.embedding = "text-embedding-3-small"
        self.client = OpenAI(base_url=config["backend_url"])
        self.chroma_client = chromadb.Client(Settings(allow_reset=True))
        self.situation_collection = self.chroma_client.create_collection(name=name)

    def get_embedding(self, text):
        """获取文本的 OpenAI 嵌入"""
        
        # 定义OpenRouter支持的嵌入模型列表，按优先级排序
        openrouter_embedding_models = [
            "openai/text-embedding-ada-002",
            "text-embedding-ada-002",
            "openai/text-embedding-3-small",
            "text-embedding-3-small"
        ]
        
        if "openrouter.ai" in str(self.client.base_url):
            # 对于OpenRouter，尝试多个模型
            for model in openrouter_embedding_models:
                try:
                    response = self.client.embeddings.create(
                        model=model, input=text
                    )
                    return response.data[0].embedding
                except Exception as e:
                    print(f"尝试模型 {model} 失败: {e}")
                    continue
            # 如果所有模型都失败，返回零向量（禁用内存功能）
            print("警告: OpenRouter不支持嵌入模型，内存功能将被禁用")
            return [0.0] * 1536  # 返回零向量
        else:
            # 对于其他提供商，使用默认模型
            try:
                response = self.client.embeddings.create(
                    model=self.embedding, input=text
                )
                return response.data[0].embedding
            except Exception as e:
                print(f"嵌入模型 {self.embedding} 失败: {e}, 返回零向量")
                return [0.0] * 1536

    def add_situations(self, situations_and_advice):
        """添加金融状况及其对应的建议。参数是元组列表 (状况, 建议)"""

        situations = []
        advice = []
        ids = []
        embeddings = []

        offset = self.situation_collection.count()

        for i, (situation, recommendation) in enumerate(situations_and_advice):
            situations.append(situation)
            advice.append(recommendation)
            ids.append(str(offset + i))
            embedding = self.get_embedding(situation)
            embeddings.append(embedding)

        # 只有当嵌入不是零向量时才添加到集合中
        if embeddings and not all(all(x == 0.0 for x in emb) for emb in embeddings):
            self.situation_collection.add(
                documents=situations,
                metadatas=[{"recommendation": rec} for rec in advice],
                embeddings=embeddings,
                ids=ids,
            )
        else:
            print("跳过添加情况到内存：嵌入功能被禁用")

    def get_memories(self, current_situation, n_matches=1):
        """使用 OpenAI 嵌入查找匹配的建议"""
        query_embedding = self.get_embedding(current_situation)

        # 检查是否是零向量（表示嵌入功能被禁用）
        if all(x == 0.0 for x in query_embedding):
            print("内存功能已禁用，返回空结果")
            return []

        # 检查集合是否为空
        if self.situation_collection.count() == 0:
            print("内存集合为空，返回空结果")
            return []

        try:
            results = self.situation_collection.query(
                query_embeddings=[query_embedding],
                n_results=n_matches,
                include=["metadatas", "documents", "distances"],
            )

            matched_results = []
            for i in range(len(results["documents"][0])):
                matched_results.append(
                    {
                        "matched_situation": results["documents"][0][i],
                        "recommendation": results["metadatas"][0][i]["recommendation"],
                        "similarity_score": 1 - results["distances"][0][i],
                    }
                )

            return matched_results
        except Exception as e:
            print(f"查询内存时出错: {e}, 返回空结果")
            return []


if __name__ == "__main__":
    # 示例用法
    matcher = FinancialSituationMemory()

    # 示例数据
    example_data = [
        (
            "高通胀率伴随利率上升和消费支出下降",
            "考虑消费必需品和公用事业等防御性行业。审查固定收益投资组合的久期。",
        ),
        (
            "科技行业表现出高波动性，机构抛售压力增加",
            "减少对高增长科技股的敞口。在现金流强劲的成熟科技公司中寻找价值机会。",
        ),
        (
            "强势美元影响新兴市场，外汇波动加剧",
            "对冲国际头寸的货币风险。考虑减少对新兴市场债务的配置。",
        ),
        (
            "市场出现行业轮动迹象，收益率上升",
            "重新平衡投资组合以维持目标配置。考虑增加对受益于利率上升的行业的敞口。",
        ),
    ]

    # 添加示例状况和建议
    matcher.add_situations(example_data)

    # 示例查询
    current_situation = """
    市场在科技行业表现出加剧的波动性，机构投资者
    减少头寸，利率上升影响成长股估值
    """

    try:
        recommendations = matcher.get_memories(current_situation, n_matches=2)

        for i, rec in enumerate(recommendations, 1):
            print(f"\n匹配 {i}:")
            print(f"相似度分数: {rec['similarity_score']:.2f}")
            print(f"匹配的状况: {rec['matched_situation']}")
            print(f"建议: {rec['recommendation']}")

    except Exception as e:
        print(f"推荐过程中出错: {str(e)}") 